'''
Created on Oct 13, 2009

@author: mkiyer
'''

from sqlalchemy import Table, Column, Integer, Float, PickleType, String, \
    MetaData, ForeignKey, create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

import math
import operator
import random

Base = declarative_base()


class Interval:
    """
    Basic feature, with required integer start and end properties.
    Also accpets optional strand as +1 or -1 (used for up/downstream queries),
    a name, and any arbitrary data is sent in on the info keyword argument

    >>> from bx.intervals.intersection import Interval

    >>> f1 = Interval(23, 36)
    >>> f2 = Interval(34, 48, value={'chr':12, 'anno':'transposon'})
    >>> f2
    Interval(34, 48, value={'anno': 'transposon', 'chr': 12})
    """
    
    def __init__(self, start, end, value=None, chrom=None, strand=None ):
        assert start <= end, "start must be less than end"
        self.start  = start
        self.end   = end
        self.value = value
        self.chrom = chrom
        self.strand = strand

    def __repr__(self):
        fstr = "Interval(%d, %d" % (self.start, self.end)
        if not self.value is None:
            fstr += ", value=" + str(self.value)
        fstr += ")"
        return fstr

    def __cmp__(self, other):
        return cmp( self.start, other.start ) or cmp( self.end, other.end )

def imax2(a, b):
    if b > a: return b
    return a

def imax3(a, b, c):
    if b > a:
        if c > b:
            return c
        return b
    if a > c:
        return a
    return c

def imin3(a, b, c):
    if b < a:
        if c < b:
            return c
        return b
    if a < c:
        return a
    return c

def imin2(a, b):
    if b < a: return b
    return a

nlog = -1.0 / math.log(0.5)


interval_table = Table('intervals', metadata,
                       Column('id', Integer, primary_key=True),
                       Column('parent_id', Integer, ForeignKey('treenodes.id')),
                       Column('data', String(50)),
                       )




class IntervalNode(Base):
    __tablename__ = 'intervals'
    
    id = Column(Integer, primary_key=True)
    priority = Column(Float)
    start = Column(Integer)
    end = Column(Integer)
    minend = Column(Integer)
    maxend = Column(Integer)
    minstart = Column(Integer)    
    value = Column(PickleType, nullable=True)

    cleft = Column(Integer, ForeignKey('intervals.id'), nullable=True)
    cright = Column(Integer, ForeignKey('intervals.id'), nullable=True)
    croot = Column(Integer, ForeignKey('intervals.id'), nullable=True)


    def __init__(self, start, end, value):        
        # Python lacks the binomial distribution, so we convert a
        # uniform into a binomial because it naturally scales with
        # tree size.  Also, python's uniform is perfect since the
        # upper limit is not inclusive, which gives us undefined here.
        self.priority = math.ceil(nlog * math.log(-1.0/(random.random() - 1)))
        self.start = start
        self.end = end
        self.maxend = end
        self.minend = end
        self.minstart = start
        self.cleft = None
        self.cright = None
        self.croot = None
        self.value = value

    def __repr__(self):
        return "<IntervalNode('%s','%s', '%s')>" % (self.start, self.end, self.value)

    def insert(self, start, end, value):
        """
        Insert a new IntervalNode into the tree of which this node is
        currently the root. The return value is the new root of the tree (which
        may or may not be this node!)
        """
        croot = self
        # If starts are the same, decide which to add interval to based on
        # end, thus maintaining sortedness relative to start/end
        decision_endpoint = start
        if start == self.start:
            decision_endpoint = end
        
        if decision_endpoint > self.start:
            # insert to cright tree
            if self.cright is not None:
                self.cright = self.cright.insert(start, end, value)
            else:
                self.cright = IntervalNode(start, end, value)
            # rebalance tree
            if self.priority < self.cright.priority:
                croot = self.rotate_left()
        else:
            # insert to cleft tree
            if self.cleft is not None:
                self.cleft = self.cleft.insert(start, end, value)
            else:
                self.cleft = IntervalNode(start, end, value)
            # rebalance tree
            if self.priority < self.cleft.priority:
                croot = self.rotate_right()
    
        croot.set_ends()
        self.cleft.croot  = croot
        self.cright.croot = croot
        return croot

    def rotate_right(self):
        croot = self.cleft
        self.cleft  = self.cleft.cright
        croot.cright = self
        self.set_ends()
        return croot

    def rotate_left(self):
        croot = self.cright
        self.cright = self.cright.cleft
        croot.cleft  = self
        self.set_ends()
        return croot

    def set_ends(self):
        if self.cright is not None and self.cleft is not None: 
            self.maxend = imax3(self.end, self.cright.maxend, self.cleft.maxend)
            self.minend = imin3(self.end, self.cright.minend, self.cleft.minend)
            self.minstart = imin3(self.start, self.cright.minstart, self.cleft.minstart)
        elif self.cright is not None:
            self.maxend = imax2(self.end, self.cright.maxend)
            self.minend = imin2(self.end, self.cright.minend)
            self.minstart = imin2(self.start, self.cright.minstart)
        elif self.cleft is not None:
            self.maxend = imax2(self.end, self.cleft.maxend)
            self.minend = imin2(self.end, self.cleft.minend)
            self.minstart = imin2(self.start, self.cleft.minstart)

    def intersect(self, start, end, sort=True):
        """
        given a start and a end, return a list of features
        falling within that range
        """
        results = []
        self._intersect(start, end, results)
        return results
        
    def _intersect(self, start, end, results):
        # Left subtree
        if self.cleft is not None and self.cleft.maxend > start:
            self.cleft._intersect(start, end, results)
        # This interval
        if (self.end > start) and (self.start < end):
            results.append(self)
        # Right subtree
        if self.cright is not None and self.start < end:
            self.cright._intersect(start, end, results)
    
    def _seek_left(self, position, results, n, max_dist):
        # we know we can bail in these 2 cases.
        if self.maxend + max_dist < position:
            return
        if self.minstart > position:
            return

        # the ordering of these 3 blocks makes it so the results are
        # ordered nearest to farest from the query position
        if self.cright is not None:
            self.cright._seek_left(position, results, n, max_dist)

        if -1 < position - self.end < max_dist:
            results.append(self)

        # TODO: can these conditionals be more stringent?
        if self.cleft is not None:
            self.cleft._seek_left(position, results, n, max_dist)

    
    def _seek_right(self, position, results, n, max_dist):
        # we know we can bail in these 2 cases.
        if self.maxend < position: 
            return
        if self.minstart - max_dist > position: 
            return

        #print "SEEK_RIGHT:",self, self.cleft, self.maxend, self.minstart, position

        # the ordering of these 3 blocks makes it so the results are
        # ordered nearest to farthest from the query position
        if self.cleft is not None: 
            self.cleft._seek_right(position, results, n, max_dist)

        if -1 < self.start - position < max_dist:
            results.append(self)

        if self.cright is not None:
            self.cright._seek_right(position, results, n, max_dist)

    
    def left(self, position, n=1, max_dist=2500):
        """
        find n features with a start > than `position`
        f: a Interval object (or anything with an `end` attribute)
        n: the number of features to return
        max_dist: the maximum distance to look before giving up.
        """
        results = []
        # use start - 1 becuase .left() assumes strictly left-of
        self._seek_left( position - 1, results, n, max_dist )
        if len(results) == n: return results
        r = results
        r.sort(key=operator.attrgetter('end'), reverse=True)
        return r[:n]

    def right(self, position, n=1, max_dist=2500):
        """
        find n features with a end < than position
        f: a Interval object (or anything with a `start` attribute)
        n: the number of features to return
        max_dist: the maximum distance to look before giving up.
        """
        results = []
        # use end + 1 becuase .right() assumes strictly right-of
        self._seek_right(position + 1, results, n, max_dist)
        if len(results) == n: return results
        r = results
        r.sort(key=operator.attrgetter('start'))
        return r[:n]

    def traverse(self, func):
        self._traverse(func)

    def _traverse(self, func):
        if self.cleft is not None: self.cleft._traverse(func)
        func(self)
        if self.cright is not None: self.cright._traverse(func)


class IntervalTree:
    """
    Data structure for performing window intersect queries on a set of 
    of possibly overlapping 1d intervals.
    
    Usage
    =====
    
    Create an empry IntervalTree
    
    >>> from bx.intervals.intersection import Interval, IntervalTree
    >>> intersecter = IntervalTree()
    
    An interval is a start and end position and a value (possibly None).
    You can add any object as an interval:
    
    >>> intersecter.insert( 0, 10, "food" )
    >>> intersecter.insert( 3, 7, dict(foo='bar') )
    
    >>> intersecter.find( 2, 5 )
    ['food', {'foo': 'bar'}]
    
    If the object has start and end attributes (like the Interval class) there
    is are some shortcuts:
    
    >>> intersecter = IntervalTree()
    >>> intersecter.insert_interval( Interval( 0, 10 ) )
    >>> intersecter.insert_interval( Interval( 3, 7 ) )
    >>> intersecter.insert_interval( Interval( 3, 40 ) )
    >>> intersecter.insert_interval( Interval( 13, 50 ) )
    
    >>> intersecter.find( 30, 50 )
    [Interval(3, 40), Interval(13, 50)]
    >>> intersecter.find( 100, 200 )
    []
    
    Before/after for intervals
    
    >>> intersecter.before_interval( Interval( 10, 20 ) )
    [Interval(3, 7)]
    >>> intersecter.before_interval( Interval( 5, 20 ) )
    []
    
    Upstream/downstread
    
    >>> intersecter.upstream_of_interval(Interval(11, 12))
    [Interval(0, 10)]
    >>> intersecter.upstream_of_interval(Interval(11, 12, strand="-"))
    [Interval(13, 50)]

    >>> intersecter.upstream_of_interval(Interval(1, 2, strand="-"), num_intervals=3)
    [Interval(3, 7), Interval(3, 40), Interval(13, 50)]

    
    """
        
    def __init__(self):
        self.root = None
    
    # ---- Position based interfaces -----------------------------------------
    
    def insert( self, start, end, value=None ):
        """
        Insert the interval [start,end) associated with value `value`.
        """
        if self.root is None:
            self.root = IntervalNode( start, end, value )
        else:
            self.root = self.root.insert( start, end, value )
        
    def find( self, start, end ):
        """
        Return a sorted list of all intervals overlapping [start,end).
        """
        return self.root.find( start, end )
    
    def before( self, position, num_intervals=1, max_dist=2500 ):
        """
        Find `num_intervals` intervals that lie before `position` and are no
        further than `max_dist` positions aways
        """
        return self.root.left( position, num_intervals, max_dist )

    def after( self, position, num_intervals=1, max_dist=2500 ):
        """
        Find `num_intervals` intervals that lie after `position` and are no
        further than `max_dist` positions aways
        """
        return self.root.right( position, num_intervals, max_dist )

    # ---- Interval-like object based interfaces -----------------------------

    def insert_interval( self, interval ):
        """
        Insert an "interval" like object (one with at least start and end
        attributes)
        """
        self.insert( interval.start, interval.end, interval )

    def before_interval( self, interval, num_intervals=1, max_dist=2500 ):
        """
        Find `num_intervals` intervals that lie completely before `interval`
        and are no further than `max_dist` positions aways
        """
        return self.root.left( interval.start, num_intervals, max_dist )

    def after_interval( self, interval, num_intervals=1, max_dist=2500 ):
        """
        Find `num_intervals` intervals that lie comletey after `interval` and
        are no further than `max_dist` positions aways
        """
        return self.root.right( interval.end, num_intervals, max_dist )

    def upstream_of_interval( self, interval, num_intervals=1, max_dist=2500 ):
        """
        Find `num_intervals` intervals that lie completely upstream of
        `interval` and are no further than `max_dist` positions aways
        """
        if interval.strand == -1 or interval.strand == "-":
            return self.root.right( interval.end, num_intervals, max_dist )
        else:
            return self.root.left( interval.start, num_intervals, max_dist )

    def downstream_of_interval( self, interval, num_intervals=1, max_dist=2500 ):
        """
        Find `num_intervals` intervals that lie completely downstream of
        `interval` and are no further than `max_dist` positions aways
        """
        if interval.strand == -1 or interval.strand == "-":
            return self.root.left( interval.start, num_intervals, max_dist )
        else:
            return self.root.right( interval.end, num_intervals, max_dist )
    
    def traverse(self, fn):
        """
        call fn for each element in the tree
        """
        return self.root.traverse(fn)



engine = create_engine('sqlite:///:memory:', echo=True)
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)
session = Session()

session.add(IntervalNode(0, 10, Interval(0, 10)))       
session.add(IntervalNode(10, 20, Interval(10, 20)))       
session.add(IntervalNode(5, 15, Interval(5, 15)))
session.add(IntervalNode(50, 100, Interval(50, 100)))        
session.add(IntervalNode(2, 6, Interval(2, 6)))

for res in session.query(IntervalNode):
    print res


